import sensor, image, time, gc, pyb
from pyb import UART, LED

sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.HQVGA)
sensor.skip_frames(time = 200)
sensor.set_auto_gain(False) # 颜色跟踪必须关闭自动增益
sensor.set_auto_whitebal(False) # 颜色跟踪必须关闭白平衡

#人为操控曝光时间，这样做会减少帧率
#正常曝光时间大概在3000~5000us
#sensor.set_auto_exposure(False, exposure_us = 15000)

uart = pyb.UART(3,9600)
greenled = pyb.LED(2)

#定义需要裁剪的ROI
ROI_Shang = (0,60,240,80)
ROI_Xia = (0,80,240,80)

#定义需要识别颜色的ROI
LeftROI_Shang = (25,15,40,60)
MiddleROI_Shang = (100,15,40,60)
RightROI_Shang = (180,14,40,60)
LeftROI_Xia = (35,20,50,60)
MiddleROI_Xia = (100,15,40,40)
RightROI_Xia = (170,20,40,55)

#左 中 右三个ROI的阈值LAB均值
LeftMean_Shang = [0, 0, 0]
MiddleMean_Shang = [0, 0, 0]
RightMean_Shang = [0, 0, 0]
LeftMean_Xia = [0, 0, 0]
MiddleMean_Xia = [0, 0, 0]
RightMean_Xia = [0, 0, 0]

#左 中 右三个ROI的中心点x值与y值
Left_x = 0
Middle_x = 0
Right_x = 0
Left_y = 0
Middle_y = 0
Right_y = 0

#红绿蓝三种颜色位置
Red_x = 0
Green_x = 0
Blue_x = 0

#颜色排序
queue = "?"

clock = time.clock()

def GetColor(Rev_Num):
    global Left_x
    global Middle_x
    global Right_x
    global Left_y
    global Middle_y
    global Right_y
    global LeftMean_Shang
    global MiddleMean_Shang
    global RightMean_Shang
    global LeftMean_Xia
    global MiddleMean_Xia
    global RightMean_Xia
    if (Rev_Num == 1):
        #获取左中右三个位置的X坐标和Y坐标
        Left_x = LeftROI_Shang[0] + (LeftROI_Shang[2] / 2)
        Middle_x = MiddleROI_Shang[0] + (MiddleROI_Shang[2] / 2)
        Right_x = RightROI_Shang[0] + (RightROI_Shang[2] / 2)
        Left_y = LeftROI_Shang[1] + (LeftROI_Shang[3] / 2)
        Middle_y = MiddleROI_Shang[1] + (MiddleROI_Shang[3] / 2)
        Right_y = RightROI_Shang[1] + (RightROI_Shang[3] / 2)
        for i in range(100):
            img = sensor.snapshot()
            img.draw_rectangle(LeftROI_Shang)
            img.draw_rectangle(MiddleROI_Shang)
            img.draw_rectangle(RightROI_Shang)
            img.draw_cross(int(Left_x), int(Left_y))
            img.draw_cross(int(Middle_x), int(Middle_y))
            img.draw_cross(int(Right_x), int(Right_y))
            Left_Stat = img.get_statistics(roi = LeftROI_Shang)
            Middle_Stat = img.get_statistics(roi = MiddleROI_Shang)
            Right_Stat = img.get_statistics(roi = RightROI_Shang)
            #获取上层左侧ROI的LAB均值
            LeftMean_Shang[0] = LeftMean_Shang[0] + (Left_Stat.l_mean() / 100)
            LeftMean_Shang[1] = LeftMean_Shang[1] + (Left_Stat.a_mean() / 100)
            LeftMean_Shang[2] = LeftMean_Shang[2] + (Left_Stat.b_mean() / 100)
            #获取上层中间ROI的LAB均值
            MiddleMean_Shang[0] = MiddleMean_Shang[0] + (Middle_Stat.l_mean() / 100)
            MiddleMean_Shang[1] = MiddleMean_Shang[1] + (Middle_Stat.a_mean() / 100)
            MiddleMean_Shang[2] = MiddleMean_Shang[2] + (Middle_Stat.b_mean() / 100)
            #获取上层右侧ROI的LAB均值
            RightMean_Shang[0] = RightMean_Shang[0] + (Right_Stat.l_mean() / 100)
            RightMean_Shang[1] = RightMean_Shang[1] + (Right_Stat.a_mean() / 100)
            RightMean_Shang[2] = RightMean_Shang[2] + (Right_Stat.b_mean() / 100)
    elif (Rev_Num == 2):
        #获取左中右三个位置的X坐标和Y坐标
        Left_x = LeftROI_Xia[0] + (LeftROI_Xia[2] / 2)
        Middle_x = MiddleROI_Xia[0] + (MiddleROI_Xia[2] / 2)
        Right_x = RightROI_Xia[0] + (RightROI_Xia[2] / 2)
        Left_y = LeftROI_Xia[1] + (LeftROI_Xia[3] / 2)
        Middle_y = MiddleROI_Xia[1] + (MiddleROI_Xia[3] / 2)
        Right_y = RightROI_Xia[1] + (RightROI_Xia[3] / 2)
        for i in range(100):
            img = sensor.snapshot()
            img.draw_rectangle(LeftROI_Xia)
            img.draw_rectangle(MiddleROI_Xia)
            img.draw_rectangle(RightROI_Xia)
            img.draw_cross(int(Left_x), int(Left_y))
            img.draw_cross(int(Middle_x), int(Middle_y))
            img.draw_cross(int(Right_x), int(Right_y))
            Left_Stat = img.get_statistics(roi = LeftROI_Xia)
            Middle_Stat = img.get_statistics(roi = MiddleROI_Xia)
            Right_Stat = img.get_statistics(roi = RightROI_Xia)
            #获取下层左侧ROI的LAB均值
            LeftMean_Xia[0] = LeftMean_Xia[0] + (Left_Stat.l_mean() / 100)
            LeftMean_Xia[1] = LeftMean_Xia[1] + (Left_Stat.a_mean() / 100)
            LeftMean_Xia[2] = LeftMean_Xia[2] + (Left_Stat.b_mean() / 100)
            #获取下层中间ROI的LAB均值
            MiddleMean_Xia[0] = MiddleMean_Xia[0] + (Middle_Stat.l_mean() / 100)
            MiddleMean_Xia[1] = MiddleMean_Xia[1] + (Middle_Stat.a_mean() / 100)
            MiddleMean_Xia[2] = MiddleMean_Xia[2] + (Middle_Stat.b_mean() / 100)
            #获取下层右侧ROI的LAB均值
            RightMean_Xia[0] = RightMean_Xia[0] + (Right_Stat.l_mean() / 100)
            RightMean_Xia[1] = RightMean_Xia[1] + (Right_Stat.a_mean() / 100)
            RightMean_Xia[2] = RightMean_Xia[2] + (Right_Stat.b_mean() / 100)

def MetchColor(Rev_Num):
    global Red_x
    global Green_x
    global Blue_x
    if (Rev_Num == 1): # 判断上层颜色排序
        #左侧颜色判断
        if (LeftMean_Shang[1] > 0 and LeftMean_Shang[2] > 0):
            Red_x = Left_x
        elif (LeftMean_Shang[1] < 0 and LeftMean_Shang[2] > 0):
            Green_x = Left_x
        elif (LeftMean_Shang[1] > 0 and LeftMean_Shang[2] < 0):
            Blue_x = Left_x
        #中间颜色判断
        if (MiddleMean_Shang[1] > 0 and MiddleMean_Shang[2] > 0):
            Red_x = Middle_x
        elif (MiddleMean_Shang[1] < 0 and MiddleMean_Shang[2] > 0):
            Green_x = Middle_x
        elif (MiddleMean_Shang[1] > 0 and MiddleMean_Shang[2] < 0):
            Blue_x = Middle_x
        #右侧颜色判断
        if (RightMean_Shang[1] > 0 and RightMean_Shang[2] > 0):
            Red_x = Right_x
        elif (RightMean_Shang[1] < 0 and RightMean_Shang[2] > 0):
            Green_x = Right_x
        elif (RightMean_Shang[1] > 0 and RightMean_Shang[2] < 0):
            Blue_x = Right_x
    elif (Rev_Num == 2):
        #左侧颜色判断
        if (LeftMean_Xia[1] > 0 and LeftMean_Xia[2] > 0):
            Red_x = Left_x
        elif (LeftMean_Xia[1] < 0 and LeftMean_Xia[2] > 0):
            Green_x = Left_x
        elif (LeftMean_Xia[1] > 0 and LeftMean_Xia[2] < 0):
            Blue_x = Left_x
        #中间颜色判断
        if (MiddleMean_Xia[1] > 0 and MiddleMean_Xia[2] > 0):
            Red_x = Middle_x
        elif (MiddleMean_Xia[1] < 0 and MiddleMean_Xia[2] > 0):
            Green_x = Middle_x
        elif (MiddleMean_Xia[1] > 0 and MiddleMean_Xia[2] < 0):
            Blue_x = Middle_x
        #右侧颜色判断
        if (RightMean_Xia[1] > 0 and RightMean_Xia[2] > 0):
            Red_x = Right_x
        elif (RightMean_Xia[1] < 0 and RightMean_Xia[2] > 0):
            Green_x = Right_x
        elif (RightMean_Xia[1] > 0 and RightMean_Xia[2] < 0):
            Blue_x = Right_x
    #打印红绿蓝的x坐标
    print ("红绿蓝X坐标：", Red_x, Green_x, Blue_x)

def ArrangeColor():
    global queue
    if (Red_x < Green_x and Green_x < Blue_x):
        queue = "123"
        return queue
    elif (Red_x < Blue_x and Blue_x < Green_x):
        queue = "132"
        return queue
    elif (Green_x < Red_x and Red_x < Blue_x):
        queue = "213"
        return queue
    elif (Green_x < Blue_x and Blue_x < Red_x):
        queue = "231"
        return queue
    elif (Blue_x < Red_x and Red_x < Green_x):
        queue = "312"
        return queue
    elif (Blue_x < Green_x and Green_x < Red_x):
        queue = "321"
    return queue

def ColorTracking(Rev_Str):
    #处理字符串,跳转获取ROI阈值并匹配颜色，然后返回排序字符串
    if (Rev_Str == "shang"):
        sensor.set_windowing(ROI_Shang)
        GetColor(1)
        #反馈ROI均值
        print ("LeftMean_Shang:", LeftMean_Shang)
        print ("MiddleMean_Shang:", MiddleMean_Shang)
        print ("RightMean_Shang:", RightMean_Shang)
        MetchColor(1)
        print ("颜色排序结果:",ArrangeColor())
    elif (Rev_Str == "xia"):
        sensor.set_windowing(ROI_Xia)
        GetColor(2)
        #反馈ROI均值
        print ("LeftMean_Xia:", LeftMean_Xia)
        print ("MiddleMean_Xia:", MiddleMean_Xia)
        print ("RightMean_Xia:", RightMean_Xia)
        MetchColor(2)
        print ("颜色排序结果:",ArrangeColor())

while(True):
    clock.tick()
    #print(clock.fps())
    Rev_str = ""
    if uart.any():
        Rev_str = uart.read()
        #print (Rev_str)
        greenled.on()
        time.sleep_ms(500)
        greenled.off()
        if (Rev_str == b'shang'):
            ColorTracking("shang")
            pyb.delay(200)
            uart.write (queue)
        elif (Rev_str == b'xia'):
            ColorTracking("xia")
            pyb.delay(200)
            uart.write (queue)
    img = sensor.snapshot()
    gc.collect()
